from typing import Optional, Literal
from pydantic import BaseModel, Field
from datetime import datetime


class ColorizeResponse(BaseModel):
    """Response model for image colorization"""
    
    success: bool = Field(description="Whether the operation was successful")
    task_id: str = Field(description="Unique task identifier")
    status: Literal["pending", "processing", "completed", "failed"] = Field(
        description="Current processing status"
    )
    original_image_url: Optional[str] = Field(
        default=None,
        description="URL of the original image"
    )
    colored_image_url: Optional[str] = Field(
        default=None,
        description="URL of the colored image"
    )
    error_message: Optional[str] = Field(
        default=None,
        description="Error message if operation failed"
    )
    processing_time: Optional[float] = Field(
        default=None,
        description="Processing time in seconds"
    )
    metadata: Optional[dict] = Field(
        default=None,
        description="Additional metadata about the processing"
    )
    created_at: datetime = Field(
        default_factory=datetime.utcnow,
        description="Task creation timestamp"
    )
    
    class Config:
        json_schema_extra = {
            "example": {
                "success": True,
                "task_id": "550e8400-e29b-41d4-a716-446655440000",
                "status": "completed",
                "original_image_url": "https://example.com/original.png",
                "colored_image_url": "https://example.com/colored.png",
                "processing_time": 15.5,
                "metadata": {
                    "style": "anime",
                    "strength": 0.7,
                    "resolution": "1024x1024"
                },
                "created_at": "2024-01-01T12:00:00Z"
            }
        }


class TaskStatusResponse(BaseModel):
    """Response model for task status check"""
    
    task_id: str = Field(description="Task identifier")
    status: Literal["pending", "processing", "completed", "failed"] = Field(
        description="Current processing status"
    )
    progress: int = Field(
        default=0,
        ge=0,
        le=100,
        description="Processing progress percentage"
    )
    result: Optional[ColorizeResponse] = Field(
        default=None,
        description="Final result if completed"
    )
    estimated_time_remaining: Optional[float] = Field(
        default=None,
        description="Estimated time remaining in seconds"
    )
    
    class Config:
        json_schema_extra = {
            "example": {
                "task_id": "550e8400-e29b-41d4-a716-446655440000",
                "status": "processing",
                "progress": 65,
                "estimated_time_remaining": 8.5
            }
        }


class HealthResponse(BaseModel):
    """Response model for health check"""
    
    status: Literal["healthy", "unhealthy"] = Field(description="Service health status")
    timestamp: datetime = Field(default_factory=datetime.utcnow)
    version: str = Field(description="API version")
    services: dict = Field(description="Status of dependent services")
    
    class Config:
        json_schema_extra = {
            "example": {
                "status": "healthy",
                "timestamp": "2024-01-01T12:00:00Z",
                "version": "1.0.0",
                "services": {
                    "database": "connected",
                    "redis": "connected",
                    "model": "loaded"
                }
            }
        }
